import numpy as np
import pandas as pd
from qtorch.strategy import Strategy

class BollingerBands(Strategy):
    """布林带策略（20日周期，2倍标准差）"""
    def __init__(self, window=20, num_std=2):
        self.window = window
        self.num_std = num_std
        
    def generate_signals(self, prices):
        rolling_mean = prices.rolling(window=self.window).mean()
        rolling_std = prices.rolling(window=self.window).std()
        
        upper_band = rolling_mean + (rolling_std * self.num_std)
        lower_band = rolling_mean - (rolling_std * self.num_std)
        
        signals = np.where(prices > upper_band, -1, 
                         np.where(prices < lower_band, 1, 0))
        return signals.astype(int)